package com.k8.cluster.loadBalanceImpl;

import com.k8.rpc.cluster.loadBalance.LoadBalance;
import com.k8.rpc.directory.Directory;
import com.k8.rpc.invoke.Invoker;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * @Author: k8
 * @CreateTime: 2025-08-17
 * @Version: 1.0
 */
public class PollBalance<T> implements LoadBalance<T> {
    private List<Invoker<T>> invokers = new ArrayList<>();
    private Directory<T> directory;
    private ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
    private boolean pushed;
    private AtomicInteger index = new AtomicInteger(0); // 初始化为0更合理

    @Override
    public void pushInvokersDirectory(Directory<T> directory) {
        if (pushed) return;
        synchronized (this) {
            if (pushed) return;
            pushed = true;
            // 尽早初始化directory，避免空指针
            this.directory = directory;
            invokers.addAll(directory.list());
            directory.addInvokerAddListener(invokers1 -> {
                try {
                    lock.writeLock().lock();
                    for (Invoker<?> tInvoker : invokers1) {
                        @SuppressWarnings("unchecked")
                        Invoker<T> invoker = (Invoker<T>) tInvoker;
                        invokers.add(invoker);
                    }
                } finally {
                    lock.writeLock().unlock();
                }
            });
        }
    }

    @Override
    public Invoker<T> nextInvoker() {
        try {
            lock.readLock().lock();
            int size = invokers.size();
            // 检查集合是否为空，避免除0
            if (size == 0) {
                throw new IllegalStateException("No available invokers for load balancing");
            }
            // 使用更简洁的原子操作实现轮询
            int currentIndex = index.getAndIncrement() % size;
            // 处理负数情况（当index溢出时）
            return invokers.get((currentIndex + size) % size);
        } finally {
            lock.readLock().unlock();
        }
    }

    @Override
    public int getSize() {
        // 处理directory可能为null的情况
        if (directory == null) {
            return invokers.size();
        }
        return directory.size();
    }
}
